# -*- coding: utf-8 -*-
import numpy as np
from algorithm.source_terms.source_base_formula import *
from algorithm.source_terms.flow_patterns.source_base import BaseSource

class AnnularFlowCalculator(BaseSource):
    def __init__(self, standard_input):
        self.standard_input = standard_input
    # def __init__(self, V_L, V_g, V_sL, V_sg, V_r, V_s, alpha_L, alpha_g, rho_L, rho_g, mu_L, mu_g, D, F_E, area, p, theta, eps, sigma):
    #     self.V_L = V_L
    #     self.V_g = V_g
    #     self.V_sL = V_sL
    #     self.V_sg = V_sg
    #     self.V_r = V_r
    #     self.V_s = V_s
    #     self.alpha_L = alpha_L
    #     self.alpha_g = alpha_g
    #     self.rho_L = rho_L
    #     self.rho_g = rho_g
    #     self.mu_L = mu_L
    #     self.mu_g = mu_g
    #     self.D = D
    #     self.F_E = F_E
    #     self.area = area
    #     self.p = p
    #     self.theta = theta
    #     self.eps = eps
    #     self.sigma = sigma
    

    @staticmethod
    def cal_delta(
        V_sg: np.ndarray,
        V_sL: np.ndarray,
        alpha_L: np.ndarray,
        D: float,  
        F_E: np.ndarray,
    ) -> np.ndarray:
        """ 计算环状流中的液膜厚度 delta

        Args:
            V_sL (np.ndarray): 液相折算速度
            V_sg (np.ndarray): 气相折算速度
            D :float 管的内径  # 改为float
            F_E :np.ndarray 夹带率

        Returns:
            np.ndarray环状流中的液膜厚度
        """
        term = (1 - alpha_L) * (V_sg + V_sL * F_E) / (V_sg + 1e-15)
        delta = (D / 2) * (1 - np.sqrt(term))
        return delta

    @staticmethod
    def cal_annular_flow_S_i(D: float, delta: np.ndarray) -> np.ndarray:
        """ 计算环状流中的界面横截面积 S_i

        Args:
            D :float 管的内径
            delta : np.ndarray 环状流中的液膜厚度

        Returns:
            np.ndarray 环状流中的界面横截面积
        """
        S_i = np.pi * (D - 2 * delta)
        return S_i

    @staticmethod
    def cal_annular_flow_S_L(D: float) ->np.ndarray:
        """ 计算环状流气相的横截面积 S_L

        Args:
            D (float): 管的内径
            
        Returns:
            np.ndarray: 液相横截面积
        """
        S_L = D * np.pi
        return S_L

    @staticmethod
    def cal_annular_flow_A_L(D: float, delta: np.ndarray) -> np.ndarray:
        """  计算环状流中的液相横截面积 A_L

        Args:
            D :float 管的内径
            delta : np.ndarray 环状流中的液膜厚度

        Returns:
            np.ndarray 环状流中的液相横截面积
        """
        A_L = np.pi * (D * delta - delta**2)
        # log.debug("函数中A_L:%s", A_L)
        return A_L

    @staticmethod
    def cal_annular_flow_A_g(D: float, delta: np.ndarray) -> np.ndarray:
        """ 计算环状流中的气相横截面积 A_g

        Args:
            D : float 管的内径
            delta : np.ndarray 环状流中的液膜厚度

        Returns:
            np.ndarray 环状流中的气相横截面积
        """
        A_g = np.pi * (D / 2 - delta) ** 2
        return A_g

    @staticmethod
    def cal_annular_flow_re(
        V_sL: np.ndarray,
        V_sg: np.ndarray,
        V_L: np.ndarray,
        V_g: np.ndarray,
        alpha_L: np.ndarray,
        mu_L: np.ndarray,
        mu_g: np.ndarray,
        rho_L: np.ndarray,
        rho_g: np.ndarray,
        D: float,
        F_E: np.ndarray,
        flag = "gas",
    ) -> np.ndarray:
        """ 计算壁面处的水力直径和雷诺数

        Args:
            V_g : np.ndarray 气相速度
            alpha_g : np.ndarray  管道中连续的气体的截面分数
            V_L : np.ndarray  液相速度
            alpha_L : np.ndarray  管道中连续的液体的截面分数
            mu_L : np.ndarray   液体黏度
            mu_g : np.ndarray   气体黏度
            rho_L : np.ndarray  液体密度
            rho_g : np.ndarray  气体密度
            area : np.ndarray  管道面面积
            D : np.ndarray  管的内径
            theta : np.ndarray  管道与水平方向的夹角
            theta_centerangle : np.ndarray  湿周圆心角
            F_E: np.ndarray 夹带率
            H_lf: np.ndarray 液膜持液率
            flag (str, optional): _description_. Defaults to "gas".
            mode (str, optional): _description_. Defaults to 1 代表是 "stratifiedflow".

        Returns:
            np.ndarray 水力直径和雷诺数
        """
        #d_h = np.zeros_like(V_g)
        re = np.zeros_like(alpha_L)
        # V_sg = cal_V_sg(V_g,alpha_g)
        # V_sL = cal_V_sg(V_L,alpha_L)
        if flag == "gas":
            delta = AnnularFlowCalculator.cal_delta(V_sg, V_sL, alpha_L, D, F_E)
            A_g = AnnularFlowCalculator.cal_annular_flow_A_g(D, delta)
            S_i = AnnularFlowCalculator.cal_annular_flow_S_i(D, delta)
            S_g = 0
            d_h = cal_d_h_g(A_g, S_i, S_g)
            re = cal_re(rho_g, d_h, V_g, mu_g)   

        elif flag == "liquid":
            delta = AnnularFlowCalculator.cal_delta(V_sg, V_sL, alpha_L, D, F_E)
            A_L = AnnularFlowCalculator.cal_annular_flow_A_L(D, delta)
            S_L = AnnularFlowCalculator.cal_annular_flow_S_L(D)
            d_h = cal_d_h_L(A_L, S_L)
            re = cal_re(rho_L, d_h, V_L, mu_L)
             
        return re

    @staticmethod
    def cal_annular_flow_wall_fric(  
        V_sL: np.ndarray,
        V_sg: np.ndarray,
        V_L: np.ndarray,
        V_g: np.ndarray,
        alpha_L: np.ndarray,
        mu_L: np.ndarray,
        mu_g: np.ndarray,
        rho_L: np.ndarray,
        rho_g: np.ndarray,
        D: float,
        F_E: np.ndarray,
        eps: float,
        flag = "gas",
    ) -> np.ndarray:
        """ 计算壁面摩擦因子

        Args:
            V_g : np.ndarray 气相速度
            alpha_g : np.ndarray 管道中连续的气体的截面分数
            V_L : np.ndarray 液相速度
            alpha_L : np.ndarray 管道中连续的液体的截面分数
            mu_L : np.ndarray 液体黏度
            mu_g : np.ndarray 气体黏度
            rho_L : np.ndarray 液体密度
            rho_g : np.ndarray 气体密度
            area : float 管道面面积
            D: float 管的内径
            sigma : np.ndarray 表面张力
            theta : np.ndarray 管道与水平方向的夹角
            theta_centerangle : np.ndarray 湿周圆心角
            eps : float 介质表面粗糙度
            flag (str, optional): _description_. Defaults to "gas".
            mode (str, optional): _description_. Defaults to 1代表是 "stratifiedflow".

        Returns:
            np.ndarray 壁面摩擦因子
        """
        re = AnnularFlowCalculator.cal_annular_flow_re(
        V_sL,
        V_sg,
        V_L,
        V_g,
        alpha_L,
        mu_L,
        mu_g,
        rho_L,
        rho_g,
        D,
        F_E,
        flag,
        )
        lam_t = AnnularFlowCalculator.cal_lam_t(re, D, eps)
        return lam_t

    @staticmethod
    def cal_lam_t(re: np.ndarray, D: float, eps: float) -> np.ndarray:
        """ 使用经验公式计算壁面摩擦因子，气液相分别和管壁的摩擦系数

        Args:
            re : np.ndarray 雷诺数
            D : float 管的内径
            eps : float 介质表面粗糙度

        Returns:
            np.ndarray 经验公式计算壁面摩擦因子
        """

        lam_t = np.zeros_like(re)

        # 雷诺数为0的情况
        mask_zero = re == 0
        lam_t[mask_zero] = 0

        # 雷诺数大于3000的情况
        mask_turbulent = re > 3000
        if np.any(mask_turbulent):
            lam_t[mask_turbulent] = AnnularFlowCalculator.calculate_lam_t(re[mask_turbulent], D, eps)

        # 雷诺数小于2000的情况
        mask_laminar = re < 2000
        lam_t[mask_laminar] = 16 / re[mask_laminar]

        # 过渡区处理
        mask_transition = (re >= 2000) & (re <= 3000)
        if np.any(mask_transition):
            # 进行插值
            lam_laminar = 16 / re[mask_transition]
            lam_turbulent = AnnularFlowCalculator.calculate_lam_t(re[mask_transition], D, eps)
            ratio = (re[mask_transition] - 2000) / 1000
            lam_t[mask_transition] = lam_laminar * (1 - ratio) + lam_turbulent * ratio
        return lam_t

    @staticmethod
    def calculate_lam_t(
        Re: np.ndarray, D: float, eps: float
    ) -> np.ndarray:
        """ 计算摩擦因子 lam_t, 公式 (1.23) 液相/气相与管壁

        Args:
            Re : np.ndarray 雷诺数
            d_h : np.ndarray 水力直径
            eps : np.ndarray

        Returns:
            np.ndarray 摩擦因子
        """
        # y0 = 1
        # error = 1
        # while error > 1e-6:
        #     y = 1 / (3.48 - 4 * np.log(2 * eps / D + 9.35 / Re / y0))
        #     error = np.sum(np.abs(y - y0))
        #     y0 = np.copy(y)
        # lam_t = y**2 #* 4
        lam_t = 0.0055 * (1 + ((2e4 * eps) / D + (1e6 / Re)) ** (1 / 3))
        return lam_t

    @staticmethod
    def cal_annular_flow_lam_i(
        re_g: np.ndarray, D: float, delta: np.ndarray
    ) -> np.ndarray:
        """ 计算环状流中的界面摩擦系数,lam_i

        Args:
            re_g : np.ndarray 雷诺数
            D : float 管的内径
            delta : np.ndarray 环状流中的液膜厚度

        Returns:
            np.ndarray 环状流中的界面摩擦系数
        """
        lam_i = 0.005 * (1 + 300 * ((1 + 17500 / re_g) * delta / D - 0.0015))
        return lam_i

    @staticmethod
    def cal_annular_flow_V_F(V_sL: np.ndarray, D: float, delta: np.ndarray, F_E: np.ndarray)->np.ndarray:
        """计算环状流的液膜流速，V_F

        Args:
            V_sL (np.ndarray): 液相折算速度
            D (float): 管的内径
            delta (np.ndarray): 环状流中的液膜厚度
            F_E : np.ndarray 夹带率

        Returns:
            np.ndarray: 环状流的液膜流速
        """        
        #V_sL = cal_V_sL(V_L, alpha_L)
        V_F = V_sL * (1 - F_E) / (4 * delta * (1 - delta / D) / D)
        return V_F

    def calculate(self):
        #self.theta_centerangle = 2 * np.pi
        self.delta = AnnularFlowCalculator.cal_delta(
            self.V_sg,
            self.V_sL,
            self.alpha_L,
            self.D,
            self.F_E,
        )
        self.S_i = AnnularFlowCalculator.cal_annular_flow_S_i(self.D, self.delta)
        self.lam_L = AnnularFlowCalculator.cal_annular_flow_wall_fric(
            self.V_sL,
            self.V_sg,
            self.V_L,
            self.V_g,   
            self.alpha_L,
            self.mu_L,
            self.mu_g,
            self.rho_L,
            self.rho_g,
            self.D,
            self.F_E,
            self.eps,
            flag="liquid",
        )
        self.A_g = AnnularFlowCalculator.cal_annular_flow_A_g(self.D, self.delta)
        self.S_g = 0
        self.d_h_g = cal_d_h_g(self.A_g, self.S_i, self.S_g)
        self.re_g = cal_re(
            self.rho_g, self.d_h_g, self.V_g, self.mu_g
        )
        self.lam_i = AnnularFlowCalculator.cal_annular_flow_lam_i(self.re_g, self.D, self.delta)
        self.rho_c = cal_rho_gD(
            self.V_sL,
            self.V_sg,
            self.rho_L,
            self.rho_g,
            self.F_E,
        )
        # self.V_F = AnnularFlowCalculator.cal_annular_flow_V_F(
        #     self.V_sL,
        #     self.D,
        #     self.delta,
        #     self.F_E
        # )

        self.S_L = AnnularFlowCalculator.cal_annular_flow_S_L(self.D)
        self.gamma_Lw = cal_separated_flow_gamma(
            self.area,
            self.lam_L,
            self.rho_L,
            self.V_L,  # self.V_F,
            self.S_L,
        )
        self.gamma_i = cal_separated_flow_gamma(
            self.area,
            self.lam_i,
            self.rho_c,
            self.V_r,
            self.S_i,
        )

        return self.gamma_gw, self.gamma_Lw, self.gamma_i

       